#pragma once

#include <memory>
#include <atomic>
#include <iostream>
#include <exception>
#include <functional>

#include "task_base.h"

template <typename TaskFunc, typename ResultFunc>
class ThreadTask : public TaskBase {
 public:
  ThreadTask(TaskFunc &&task, ResultFunc &&result)
      : task_(std::move(task)), result_(std::move(result)) {}

  void operator()() override {
    using Ret = std::invoke_result_t<TaskFunc>;
    std::exception_ptr ex_ptr;
    if constexpr (std::is_same_v<Ret, void>) {
      try {
        task_();
      } catch (...) {
        ex_ptr = std::current_exception();
      }

      try {
        result_(ex_ptr);
      } catch (const std::exception &ex) {
        std::cerr << "Uncaught exception in ThreadTask callback: " << ex.what()
                  << std::endl;
      } catch (...) {
        std::cerr << "Uncaught unknow exception in ThreadTask callback"
                  << std::endl;
      }
    } else {
      Ret r;
      try {
        r = task_();
      } catch (...) {
        ex_ptr = std::current_exception();
      }

      try {
        result_(std::move(r), ex_ptr);
      } catch (const std::exception &ex) {
        std::cerr << "Uncaught exception in ThreadTask callback: " << ex.what()
                  << std::endl;
      } catch (...) {
        std::cerr << "Uncaught unknow exception in ThreadTask callback"
                  << std::endl;
      }
    }
  }

  template <typename _TaskFunc = TaskFunc, typename _ResultFunc = ResultFunc>
  static TaskBase *MakeTask(_TaskFunc &&task, _ResultFunc &&result) {
    return new ThreadTask<_TaskFunc, _ResultFunc>(std::move(task),
                                                  std::move(result));
  }

 private:
  TaskFunc task_;
  ResultFunc result_;
};

class ThreadTaskIteratorBase
    : public std::enable_shared_from_this<ThreadTaskIteratorBase> {
 public:
  virtual ~ThreadTaskIteratorBase() {}
  virtual TaskBase *operator*() = 0;
};

template <typename It, typename TaskFunc, typename ResultFunc>
class ThreadTaskIterator : public ThreadTaskIteratorBase {
 public:
  ThreadTaskIterator(It &&iter, TaskFunc &&task, ResultFunc &&result,
                     size_t total)
      : iter_(std::forward<It>(iter)),
        task_(std::forward<TaskFunc>(task)),
        result_(std::forward<ResultFunc>(result)),
        total_(total) {}

  TaskBase *operator*() override {
    using TaskIteratorSptr = std::shared_ptr<ThreadTaskIterator>;
    TaskIteratorSptr self =
        std::static_pointer_cast<ThreadTaskIterator>(shared_from_this());

    auto tf = [iter_ = iter_, self] { self->task_(iter_); };
    auto rf = [self](std::exception_ptr ex_ptr) mutable {
      if (self->count_ < self->total_) {
        if (ex_ptr) {
          self->result_(ex_ptr);
        } else if (++self->count_ == self->total_) {
          self->result_(ex_ptr);
        }
      }
    };
    ++iter_;
    return ThreadTask<TaskFunc, ResultFunc>::template MakeTask(std::move(tf),
                                                               std::move(rf));
  }

  It iter_;
  TaskFunc task_;
  ResultFunc result_;
  size_t total_{0};
  std::atomic_int32_t count_{0};
};

class ThreadTaskIteratorWrapper {
 public:
  ThreadTaskIteratorWrapper(std::shared_ptr<ThreadTaskIteratorBase> &&iter)
      : iter_(std::move(iter)) {}

  ThreadTaskIteratorWrapper &operator++() { return *this; }
  ThreadTaskIteratorWrapper operator++(int) {
    ThreadTaskIteratorWrapper tmp = *this;
    return tmp;
  }
  TaskBase *operator*() const { return *(*iter_); }

  std::shared_ptr<ThreadTaskIteratorBase> iter_;
};